536E - Tavas on the Path - CodeForces Solution

#include<bits/stdc++.h>
using namespace std;
const int maxn=100000+10;
struct yal{
	int u,v,w;
	int getad(int fu){
		return (fu^v^u);
	}
};

vector<int>adj[maxn];
yal alled[maxn];
int n,q,val[maxn],ph[maxn],p[maxn],sz[maxn],timea=0;
pair<int,int>stf[maxn];
int inf=1e9+5;

int kaf=(1<<17);
struct javab{
	int res,av,dov,len;
	javab(){
		res=av=dov=len=0;
	}
}alak;

javab mergej(javab a,javab b){
	javab ret;
	ret.len=a.len+b.len;
	ret.res=a.res+b.res;
	if(a.av==a.len){
		ret.av=a.av+b.av;
	}
	else{
		ret.av=a.av;
	}
	if(b.dov==b.len){
		ret.dov=b.dov+a.dov;
	}
	else{
		ret.dov=b.dov;
	}
	if(a.dov!=a.len&&b.av!=b.len){
		ret.res+=val[a.dov+b.av];
	}
	return ret;
}

struct segment{
	struct node{
		vector<int>v;
		vector<javab>allr;
		void build(){
			v.push_back(inf);
			sort(v.begin(),v.end());
			v.resize(unique(v.begin(),v.end())-v.begin());
			allr.resize(v.size());
			allr.back().res=allr.back().len=allr.back().av=allr.back().dov=0;
		}
	};
	node seg[(1<<18)];
	void ins(int i,int w){
		if(i==0){
			return ;
		}
		seg[i].v.push_back(w);
		ins((i>>1),w);
	}

	void calc(int i){
		for(int j=0;j<(int)seg[i].v.size();j++){
			int fl=lower_bound(seg[(i<<1)].v.begin(),seg[(i<<1)].v.end(),seg[i].v[j])-seg[(i<<1)].v.begin();
			int fr=lower_bound(seg[(i<<1)^1].v.begin(),seg[(i<<1)^1].v.end(),seg[i].v[j])-seg[(i<<1)^1].v.begin();
			seg[i].allr[j]=mergej(seg[(i<<1)].allr[fl],seg[(i<<1)^1].allr[fr]);
		}
	}

	void builda(){
		for(int i=0;i<(1<<18);i++){
			seg[i].build();
		}
		for(int i=(1<<18)-1;i>=1;i--){
			if(i>=kaf){
				if(seg[i].v.size()>1){
					seg[i].allr[0].av=seg[i].allr[0].len=seg[i].allr[0].dov=1;
					seg[i].allr[1].len=1;
				}
			}
			else{
				calc(i);
			}
		}
	}
	javab pors(int i,int l,int r,int tl,int tr,int w){
		if(l>r||l>tr||r<tl||tl>tr){
			return alak;
		}
		if(l>=tl&&r<=tr){
			int p=lower_bound(seg[i].v.begin(),seg[i].v.end(),w)-seg[i].v.begin();
			//cout<<"tof: "<<l<<" "<<r<<" "<<seg[i].allr[p].res<<" "<<seg[i].allr[p].len<<"\n";
			return seg[i].allr[p];
		}
		int m=(l+r)>>1;
		return mergej(pors((i<<1),l,m,tl,tr,w),pors((i<<1)^1,m+1,r,tl,tr,w));
	}
}seg;

bool zird(int u,int v){
	if(stf[v].first<=stf[u].first&&stf[v].second>=stf[u].second){
		return 1;
	}
	return 0;
}

void calseg(){
	for(int i=1;i<=n;i++){
		if(ph[i]!=i){
			seg.ins(stf[i].first+kaf,alled[p[i]].w);
		}
	}
	seg.builda();
}

bool cmp(pair<int,int>a,pair<int,int>b){
	return sz[a.first]>sz[b.first];
}

void pre(int u=1,int par=-1){
	sz[u]=1;
	vector<pair<int,int>>fake;
	for(auto xx:adj[u]){
		int x=alled[xx].getad(u);
		if(x==par){
			continue;
		}
		fake.push_back(make_pair(x,xx));
		pre(x,u);
		sz[u]+=sz[x];
		p[x]=xx;
	}
	sort(fake.begin(),fake.end(),cmp);
	adj[u].clear();
	for(int i=0;i<(int)fake.size();i++){
		adj[u].push_back(fake[i].second);
	}
}

void dfs(int u=1,int bl=1){
	timea++;
	stf[u].first=timea;
	ph[u]=bl;
	if(adj[u].size()==0){
		stf[u].second=timea;
		return ;
	}
	dfs(alled[adj[u][0]].getad(u),bl);
	for(int i=1;i<(int)adj[u].size();i++){
		int x=alled[adj[u][i]].getad(u);
		dfs(x,x);
	}
	stf[u].second=timea;
}

int getlca(int u,int v){
	if(zird(u,v)){
		return v;
	}
	if(zird(v,u)){
		return u;
	}
	int fu=u,fv=v;
	while(zird(fv,u)==0){
		u=alled[p[ph[u]]].getad(ph[u]);
	}
	while(zird(fu,v)==0){
		v=alled[p[ph[v]]].getad(ph[v]);
	}
	if(zird(u,v)){
		return u;
	}
	return v;
}

javab tabala(int u,int had,int w){
	//cout<<"ha: "<<u<<" "<<had<<" "<<w<<"\n";
	if(zird(had,ph[u])==0){
		javab jj;
		if(alled[p[ph[u]]].w>=w){
			jj.len=jj.av=jj.dov=1;
		}
		else{
			jj.len=1;
		}
		jj=mergej(jj,seg.pors(1,0,kaf-1,stf[ph[u]].first+1,stf[u].first,w));
		return mergej(tabala(alled[p[ph[u]]].getad(ph[u]),had,w),jj);
	}
	else{
		return seg.pors(1,0,kaf-1,stf[had].first+1,stf[u].first,w);
	}
}

void solve(int u,int v,int w){
	int uv=getlca(u,v);
//	//cout<<"ha: "<<u<<" "<<v<<" "<<uv<<"\n";
//	return ;
	javab av=tabala(u,uv,w);
	javab dov=tabala(v,uv,w);
	swap(av.av,av.dov);
	javab asl=mergej(av,dov);
	//cout<<av.len<<" "<<av.av<<" "<<av.dov<<"\n";
	//cout<<dov.len<<" "<<dov.av<<" "<<dov.dov<<"\n";
	////cout<<asl.len<<" "<<asl.av<<" "<<asl.dov<<"\n";
	if(asl.len!=asl.av){
		asl.res+=val[asl.av]+val[asl.dov];
	}
	else{
		asl.res+=val[asl.len];
	}
	cout<<asl.res<<"\n";
}

int main(){
	ios::sync_with_stdio(0);
	cin.tie(0);
	cout.tie(0);
	cin>>n>>q;
	for(int i=1;i<n;i++){
		cin>>val[i];
	}
	for(int i=0;i<n-1;i++){
		cin>>alled[i].u>>alled[i].v>>alled[i].w;
		adj[alled[i].u].push_back(i);
		adj[alled[i].v].push_back(i);
	}
	p[1]=-1;
	pre();
	dfs();
	calseg();
	for(int i=0;i<q;i++){
		int u,v,w;
		cin>>u>>v>>w;
		solve(u,v,w);
	}
}